
#include <iostream>
#include "FEM2DVector.h"
#include "TimeDiscrete.h"
#define PI 3.1415926
double nu(CPoint p, double t)
{
	return 1.0;
}
double u1_t0(CPoint p)
{
	double x = p.getX();
	double y = p.getY();
	return pow(x, 2)*pow(y, 2) + exp(-y);
}
double u2_t0(CPoint p)
{
	double x = p.getX();
	double y = p.getY();
	return -2.0 / 3.0*x*pow(y, 3) + 2.0 - PI * sin(PI*x);
}
double p_t0(CPoint p)
{
	double x = p.getX();
	double y = p.getY();
	return -(2-PI*sin(PI*x))*cos(2*PI*y);
}
double u1_left(CPoint p, double t)
{
	double x = p.getX();
	double y = p.getY();
	return exp(-y)*cos(2 * PI*t);
}
double u1_right(CPoint p, double t)
{
	double x = p.getX();
	double y = p.getY();
	return (pow(y, 2) + exp(-y))*cos(2 * PI*t);
}
double u1_bottom(CPoint p, double t)
{
	double x = p.getX();
	double y = p.getY();
	return (1.0 / 16.0*pow(x, 2) + exp(0.25))*cos(2 * PI*t);
}

double u1_up(CPoint p, double t)
{
	double x = p.getX();
	double y = p.getY();
	return cos(2 * PI*t);
}
double u2_left(CPoint p, double t)
{
	double x = p.getX();
	double y = p.getY();
	return 2 * cos(2 * PI*t);
}
double u2_right(CPoint p, double t)
{
	double x = p.getX();
	double y = p.getY();
	return (-2.0 / 3.0*pow(y, 3) + 2)*cos(2 * PI*t);
}
double u2_bottom(CPoint p, double t)
{
	double x = p.getX();
	double y = p.getY();
	return (1.0 / 96.0*x + 2 - PI * sin(PI*x))*cos(2 * PI*t);
}
double u2_up(CPoint p, double t)
{
	double x = p.getX();
	double y = p.getY();
	return (2 - PI * sin(PI*x))*cos(2 * PI*t);
}
double f1(CPoint p, double t)
{
	double x = p.getX();
	double y = p.getY();
	return -2 * PI*(pow(x, 2)*pow(y, 2) + exp(-y))*sin(2 * PI*t)+
		(-2*nu(p,t)*pow(x,2)-2*nu(p,t)*pow(y,2)-nu(p,t)*exp(-y)+pow(PI,2)*cos(PI*x)*cos(2*PI*y))*cos(2*PI*t);
}

double f2(CPoint p, double t)
{
	double x = p.getX();
	double y = p.getY();
	return -2 * PI*(-2.0 / 3.0*x*pow(y, 3) + 2.0 - PI * sin(PI*x))*sin(2 * PI*t) +
		(4 * nu(p, t)*x*y - nu(p, t)*pow(PI, 3)*sin(PI*x) +
			2 * PI*(2 - PI * sin(PI*x))*sin(2 * PI*y))*cos(2 * PI*t);
}
double fun_negative_one(CPoint, double)
{
	return -1.0;
}
double fun_one(CPoint, double)
{
	return 1.0;
}

int main()
{
	int unknow_size = 3;
	CFEM2DVector fem(unknow_size);

	//read mesh from file
	//the mesh file contains the P,T,Pb,Tb,boundary nodes,boundary edges, and the type of the element
	fem.readMesh("..\\mesh\\cha8_exa1.neu");

	fem.Init();	

	//assign boundary nodes value by the boundary index defined in the .neu mesh file	
	fem.assignBoundNodesValue(0, 0, Dirichlet, u1_left);
	fem.assignBoundNodesValue(1, 0, Dirichlet, u1_bottom);
	fem.assignBoundNodesValue(2, 0, Dirichlet, u1_right);
	fem.assignBoundNodesValue(3, 0, Dirichlet, u1_up);

	fem.assignBoundNodesValue(0, 1, Dirichlet, u2_left);
	fem.assignBoundNodesValue(1, 1, Dirichlet, u2_bottom);
	fem.assignBoundNodesValue(2, 1, Dirichlet, u2_right);
	fem.assignBoundNodesValue(3, 1, Dirichlet, u2_up);

	//time discrete
	double time_left = 0.0, time_right = 1.0;
	int number_of_time_step = 10;
	double theta = 0.5;
	CTimeDiscrete time(time_left, time_right, number_of_time_step, theta);

	CBasisFunction u_basis_trial(TWO_DIM_QUADRATIC);
	CBasisFunction u_basis_test(TWO_DIM_QUADRATIC);

	CBasisFunction p_basis_trial(TWO_DIM_LINEAR);
	CBasisFunction p_basis_test(TWO_DIM_LINEAR);

	pair<int, int> matrix_size_u{ fem.getPbTest().size(),fem.getPbTrial().size() };
	pair<int, int> matrix_size_p{ fem.getPbTest().size(),fem.getP().size() };

	//assemble matrix A
	SpMat A1 = fem.assembleMatrix2D(nu, 0, u_basis_trial, 1, 0, fem.getTbTrial(), u_basis_test, 1, 0, fem.getTbTest(), matrix_size_u);
	SpMat A2 = fem.assembleMatrix2D(nu, 0, u_basis_trial, 0, 1, fem.getTbTrial(), u_basis_test, 0, 1, fem.getTbTest(), matrix_size_u);
	SpMat A3 = fem.assembleMatrix2D(nu, 0, u_basis_trial, 1, 0, fem.getTbTrial(), u_basis_test, 0, 1, fem.getTbTest(), matrix_size_u);
	SpMat A5 = fem.assembleMatrix2D(fun_negative_one, 0, p_basis_trial, 0, 0, fem.getT(), u_basis_test, 1, 0, fem.getTbTest(), matrix_size_p);
	SpMat A6 = fem.assembleMatrix2D(fun_negative_one, 0, p_basis_trial, 0, 0, fem.getT(), u_basis_test, 0, 1, fem.getTbTest(), matrix_size_p);
	SpMat zero_1(fem.getP().size(), fem.getP().size());
	zero_1.setZero();
	SpMat A = fem.matrix_vstack(fem.matrix_hstack(2 * A1 + A2, A3, A5), fem.matrix_hstack(A3.transpose(), 2 * A2 + A1, A6), fem.matrix_hstack(A5.transpose(), A6.transpose(), zero_1));
	
	//assemble mass matrix M
	SpMat Me = fem.assembleMatrix2D(fun_one, 0, u_basis_trial, 0, 0, fem.getTbTrial(), u_basis_test, 0, 0, fem.getTbTest(), matrix_size_u);
	SpMat zero_2(matrix_size_p.first, matrix_size_p.second);
	SpMat zero_3(matrix_size_u.first, matrix_size_u.second);
	
	SpMat M = fem.matrix_vstack(fem.matrix_hstack(Me, zero_3, zero_2), fem.matrix_hstack(zero_3, Me, zero_2), fem.matrix_hstack(zero_2.transpose(), zero_2.transpose(), zero_1));

	//set the initial value
	Rsv u1_0 = fem.generate_Initial(u1_t0, fem.getPbTest());
	Rsv u2_0 = fem.generate_Initial(u2_t0, fem.getPbTest());
	Rsv p_0 = fem.generate_Initial(p_t0, fem.getP());
	Rsv x0 = fem.vector_stack(u1_0, u2_0, p_0);
	Rsv x;

	Rsv vec_zero(matrix_size_p.second);
	vec_zero.setZero();
	for (int m = 0; m < time.getTimeStepNumber(); ++m)
	{

		double dt = time.getTimeStep(m);
		double tm = time.getTime(m);
		double tmp1 = time.getTime(m + 1);
		cout << "Time step:\t" << m << ", time:\t" << tmp1 << endl;

		Rsv b1mp1 = fem.assembleVector2D(f1, tmp1, u_basis_test, 0, 0,fem.getTbTest());
		Rsv b1m = fem.assembleVector2D(f1, tm, u_basis_test, 0, 0, fem.getTbTest());
		Rsv b2mp1 = fem.assembleVector2D(f2, tmp1, u_basis_test, 0, 0, fem.getTbTest());
		Rsv b2m = fem.assembleVector2D(f2, tm, u_basis_test, 0, 0, fem.getTbTest());

		Rsv bm = fem.vector_stack(b1m, b2m, vec_zero);
		Rsv bmp1 = fem.vector_stack(b1mp1, b2mp1, vec_zero);

		//dt may not be the same for all time step
		SpMat Atilde = M / dt + time.getTheta()*A;
		SpMat Afixed = M / dt - (1 - time.getTheta())*A;

		Rsv btilde = time.getTheta()*bmp1 + (1 - time.getTheta())*bm + Afixed * x0;

		//treat boundary condition
		fem.treatBoundaryConditions(fun_one, tmp1, Atilde, btilde, vector<unsigned int>{fem.getPbTrial().size(), fem.getPbTrial().size(), fem.getP().size()});

		//impose a value for p
		fem.imposeFixValue(1 + 2 * matrix_size_u.first, 0, Atilde, btilde);

		//solve the linear equations
		x = fem.solveLinearEqua(DIREC_PardisoLU, Atilde, btilde);

		x0 = x;

		string filename = "u_" + to_string(m) + ".vtk";
		fem.export2VTK(filename, x);

		Rsv x_p(fem.getP().size());
		for (int i = 0; i < x_p.rows(); ++i)
		{
			x_p(i) = x(i + 2 * fem.getPbTest().size());
		}
		filename = "p_" + to_string(m) + ".vtk";
		fem.export2VTK(filename, x_p, fem.getP(), fem.getT());
	}
}


